#include <mpi.h>
#include <stdio.h>
#include "utils.h"

#define N 4  // 矩阵的大小，为了方便观察结果定义的较小

int main(int argc, char** argv) {
    int rank, size;
    int *A, *B, *C;

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    if (size != 2) {
        fprintf(stderr, "This program is designed to run with 2 processes.\n");
        MPI_Abort(MPI_COMM_WORLD, 1);
    }

    //! 分配内存并初始化矩阵
    A = (int*)malloc(N * N * sizeof(int));
    B = (int*)malloc(N * N * sizeof(int));
    C = (int*)malloc(N * N * sizeof(int));

    for (int i = 0; i < N * N; i++) {
        A[i] = rank + 1;
        B[i] = rank + 1;
        C[i] = 0;
    }

    printf("rank %d ");
    for (int i = 0; i < N*N; i++) {
        printf("%d ");
    }
    printf("\n");


    //! 创建窗口
    /**
     * @brief           创建窗口对象
     * 
     * @param base      窗口的起始地址(该地址被别的进程通过RMA访问，也可被成为共享区域大小)
     * @param size      窗口大小(共享区域大小)，以字节为单位
     * @param disp_unit 内存分配单元的大小（以字节为单位）。这个参数定义了内存地址的粒度。通常设置为
     *                  sizeof(datatype)，其中 datatype 是你将要使用的数据类型的大小, 通常设置为1就好.
     * @param info      传递给 MPI 的额外信息。通常设置为 MPI_INFO_NULL，表示没有额外的信息。
     * @param comm      通信子, RMA只能在组内通信子间通信
     * @param win       创建的窗口对象
     * 
     * @return OMPI_DECLSPEC 
     */
    // TODO: 这里使用的是 sizeof(int)
    MPI_Win_create(A, N * N * sizeof(int), sizeof(int), MPI_INFO_NULL, MPI_COMM_WORLD, &winA);
    MPI_Win_create(B, N * N * sizeof(int), sizeof(int), MPI_INFO_NULL, MPI_COMM_WORLD, &winB);
    MPI_Win_create(C, N * N * sizeof(int), sizeof(int), MPI_INFO_NULL, MPI_COMM_WORLD, &winC);


    //! RMA数据读取
    MPI_Win_fence(0, winA);            // 同步开启区域，表示可以开始RMA操作了
    MPI_Win_fence(0, winB);
    MPI_Win_fence(0, winC);
    if (rank == 0) {
        /**
         *! @brief   进程0从进程1获取数据并相加
         * 
         * @param origin_addr       源进程的消息缓冲区地址
         * @param origin_count      源进程消息缓冲区中的数据个数
         * @param origin_datatype   数据类型
         * @param target_rank       目标进程ID
         * @param target_disp       目标进程地址相对base的偏移
         * @param target_count      数据量大小
         * @param target_datatype   数据类型
         * @param win               窗口句柄
         * 
         * @return OMPI_DECLSPEC 
         */
        MPI_Get(A, N * N, MPI_INT, 1, 0, N * N, MPI_INT, winA);
        MPI_Get(B, N * N, MPI_INT, 1, 0, N * N, MPI_INT, winB);
        for (int i = 0; i < N * N; i++) {
            C[i] = A[i] + B[i];
        }
    }
    // MPI_Win_fence(0, winA);            // 同步开启区域，表示可以开始RMA操作了
    // MPI_Win_fence(0, winB);
    // MPI_Win_fence(0, winC);


    //! 释放窗口和内存
    MPI_Win_free(&winA);
    MPI_Win_free(&winB);
    MPI_Win_free(&winC);
    free(A);
    free(B);
    free(C);


    // 结束 MPI 环境
    MPI_Finalize();

    return 0;
}